{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# PREDOMINATELY TAKEN FROM https://github.com/SyneRBI/SIRF-Exercises/blob/master/notebooks/Synergistic/BrainWeb.ipynb"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install brainweb\n",
    "!pip install pyellipsoid"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import brainweb\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "url = r\"http://brainweb.bic.mni.mcgill.ca/cgi/brainweb1?do_download_alias=subject04_crisp&format_value=raw_short&zip_value=gnuzip&download_for_real=%5BStart+download%21%5D\"\n",
    "fname = \"subject_04.bin.gz\"\n",
    "files = brainweb.get_file(fname, url, \"brainweb_3D/\")\n",
    "fname = \"brainweb_3D/\" + fname\n",
    "data = brainweb.load_file(fname)\n",
    "brainweb.seed(1337)\n",
    "# You have to change the brainweb source to allow the outres of MR (1,1,1) mm\n",
    "# see https://github.com/casperdcl/brainweb/issues/12 we chose to take ceil\n",
    "vol = brainweb.get_mmr_fromfile(\n",
    "    fname,\n",
    "    petNoise=1, t1Noise=0.75, t2Noise=0.75,\n",
    "    petSigma=1, t1Sigma=1, t2Sigma=1, outres = \"MR\")\n",
    "vol_amyl = brainweb.get_mmr_fromfile(\n",
    "    fname,\n",
    "    petNoise=1, t1Noise=0.75, t2Noise=0.75,\n",
    "    petSigma=1, t1Sigma=1, t2Sigma=1,\n",
    "    PetClass=brainweb.Amyloid, outres = \"MR\")\n",
    "\n",
    "FDG_arr  = vol['PET']\n",
    "amyl_arr = vol_amyl['PET']\n",
    "uMap_arr = vol['uMap']\n",
    "T1_arr   = vol['T1']\n",
    "T2_arr   = vol['T2']\n",
    "\n",
    "for vol_path in [\"brainweb_3D/subject_04.FDG.MR.npz\",\"brainweb_3D/subject_04.Amyloid.MR.npz\"]:\n",
    "    vol = np.load(vol_path)\n",
    "    vol_new = {}\n",
    "    for vol_file in vol.files:\n",
    "        vol_new[vol_file] = vol[vol_file]\n",
    "    for vol_type in [\"PET\",\"uMap\",\"T1\",\"T2\"]:\n",
    "        vol_new[vol_type] = vol_new[vol_type][38:218, 231:-231, 231:-231]\n",
    "    np.save(vol_path[:-3]+\"ground_truth\", vol_new)\n",
    "\n",
    "def subplot_(idx,vol,title,clims=None,cmap=\"viridis\"):\n",
    "    plt.subplot(*idx)\n",
    "    plt.imshow(vol,cmap=cmap)\n",
    "    if not clims is None:\n",
    "        plt.clim(clims)\n",
    "    plt.colorbar()\n",
    "    plt.title(title)\n",
    "    plt.axis(\"off\")\n",
    "\n",
    "fdg = np.load(\"brainweb_3D/subject_04.FDG.MR.ground_truth.npy\", allow_pickle=True).item()\n",
    "amyl = np.load(\"brainweb_3D/subject_04.Amyloid.MR.ground_truth.npy\", allow_pickle=True).item()\n",
    "plt.figure();\n",
    "slice_show = 100\n",
    "subplot_([2,3,1],fdg[\"PET\"][slice_show],'FDG',cmap=\"hot\")\n",
    "subplot_([2,3,2],amyl[\"PET\"][slice_show],'Amyloid',cmap=\"hot\")\n",
    "subplot_([2,3,3],fdg[\"uMap\"][slice_show],'uMap',cmap=\"bone\")\n",
    "subplot_([2,3,4],fdg[\"T1\"][slice_show],'T1',cmap=\"Greys_r\")\n",
    "subplot_([2,3,5],fdg[\"T2\"][slice_show],'T2',cmap=\"Greys_r\")\n",
    "print(\"The resolution is: \", vol[\"res\"], \"mm. The dimensions are: \", FDG_arr.shape, \"voxels.\")\n",
    "print(\"Keys for the dicts: \", fdg.keys())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "import numpy as np\n",
    "from skimage.filters import gaussian\n",
    "import cupy as cp\n",
    "import cupyx.scipy.ndimage as ndi\n",
    "from pyellipsoid import drawing\n",
    "import torch\n",
    "\n",
    "for path in [\"brainweb_3D/subject_04.FDG.MR.ground_truth.npy\", \"brainweb_3D/subject_04.Amyloid.MR.ground_truth.npy\"]:\n",
    "    dict_gt = np.load(path, allow_pickle=True).item()\n",
    "    pet_vol = dict_gt[\"PET\"]\n",
    "    seg = cp.zeros_like(pet_vol)\n",
    "    foreground, background = 1, 0\n",
    "    seg[dict_gt[\"uMap\"] <= 0] = background\n",
    "    seg[dict_gt[\"uMap\"] > 0] = foreground\n",
    "    seg = ndi.binary_opening(seg, structure=cp.ones((3,3,3)))\n",
    "    seg = ndi.binary_closing(seg, structure=cp.ones((20,20,20)), border_value=1)\n",
    "    mask = seg.copy()\n",
    "    for _ in range(15):\n",
    "        mask = ndi.binary_erosion(mask, border_value=1)\n",
    "    mask_np = mask.get()\n",
    "\n",
    "    tumours_dict = {}\n",
    "    # Define an image shape, axis order is: Z, Y, X\n",
    "    image_shape = mask_np.shape\n",
    "\n",
    "    # Define an image shape, axis order is: Z, Y, X\n",
    "    image_shape = mask_np.shape\n",
    "    # Define an ellipsoid, axis order is: X, Y, Z\n",
    "    ell_center = [150,182,110]\n",
    "    ell_radii = (8, 19, 9)\n",
    "    ell_angles = (1231, 12312, 23432) # Order of rotations is X, Y, Z\n",
    "    # Draw a 3D binary image containing the ellipsoid\n",
    "    tumours_dict[\"tumour_mask_0\"] = drawing.make_ellipsoid_image(image_shape, ell_center, ell_radii, ell_angles)\n",
    "\n",
    "    # Define an ellipsoid, axis order is: X, Y, Z\n",
    "    ell_center = [94,62,95]\n",
    "    ell_radii = (15, 5, 10)\n",
    "    ell_angles = (213,40,60) # Order of rotations is X, Y, Z\n",
    "    # Draw a 3D binary image containing the ellipsoid\n",
    "    tumours_dict[\"tumour_mask_1\"] = drawing.make_ellipsoid_image(image_shape, ell_center, ell_radii, ell_angles)\n",
    "\n",
    "    # Define an image shape, axis order is: Z, Y, X\n",
    "    image_shape = mask_np.shape\n",
    "    # Define an ellipsoid, axis order is: X, Y, Z\n",
    "    ell_center = [102,162,25]\n",
    "    ell_radii = (12, 14, 8)\n",
    "    ell_angles = (0, 0, 0) # Order of rotations is X, Y, Z\n",
    "    # Draw a 3D binary image containing the ellipsoid\n",
    "    tumours_dict[\"tumour_mask_2\"] = drawing.make_ellipsoid_image(image_shape, ell_center, ell_radii, ell_angles)\n",
    "\n",
    "    # Define an image shape, axis order is: Z, Y, X\n",
    "    image_shape = mask_np.shape\n",
    "    # Define an ellipsoid, axis order is: X, Y, Z\n",
    "    ell_center = [93,122,62]\n",
    "    ell_radii = (16, 8, 12)\n",
    "    ell_angles = (123,21,53) # Order of rotations is X, Y, Z\n",
    "    # Draw a 3D binary image containing the ellipsoid=\n",
    "    tumours_dict[\"tumour_mask_3\"] = drawing.make_ellipsoid_image(image_shape, ell_center, ell_radii, ell_angles)\n",
    "\n",
    "    # Define an image shape, axis order is: Z, Y, X\n",
    "    image_shape = mask_np.shape\n",
    "    # Define an ellipsoid, axis order is: X, Y, Z\n",
    "    ell_center = [142,112,142]\n",
    "    ell_radii = (9, 10, 6)\n",
    "    ell_angles = (1231, 12312, 23432) # Order of rotations is X, Y, Z\n",
    "    # Draw a 3D binary image containing the ellipsoid\n",
    "    tumours_dict[\"tumour_mask_4\"] = drawing.make_ellipsoid_image(image_shape, ell_center, ell_radii, ell_angles)\n",
    "\n",
    "\n",
    "    intensities = [0.78,.76,1.06,.89,.92]\n",
    "    tumours = np.zeros_like(pet_vol)\n",
    "    backgrounds = np.zeros_like(pet_vol)\n",
    "    for i in range(5):\n",
    "        tumour_mask = tumours_dict[\"tumour_mask_\"+str(i)].astype(float)\n",
    "        outer = cp.array(tumour_mask)\n",
    "        for _ in range(15):\n",
    "            outer = ndi.binary_dilation(outer,border_value=0)\n",
    "        inner = cp.array(tumour_mask)\n",
    "        for _ in range(8):\n",
    "            inner = ndi.binary_dilation(inner, border_value=0)\n",
    "        background_mask = mask_np*(outer.get().astype(float)-inner.get().astype(float))\n",
    "        tumours_dict[\"background_mask_\"+str(i)] = background_mask\n",
    "        tumours += gaussian(tumour_mask, sigma=1)*intensities[i]*pet_vol.max()\n",
    "        backgrounds += background_mask\n",
    "\n",
    "    pet_tumours = np.max((pet_vol, tumours), axis=0)\n",
    "\n",
    "    vol = np.load(path, allow_pickle=True).item()\n",
    "    vol[\"tumours_dict\"] = tumours_dict\n",
    "    vol[\"PET_tumours\"] = pet_tumours\n",
    "    np.save(path[:-3]+\"tumours\", vol)\n",
    "\n",
    "    \"\"\" for i in range(len(pet_vol)):#unique_axial_coords:\n",
    "        plt.figure(figsize=(15,5));\n",
    "        slice_show = i\n",
    "        subplot_([1,3,1], pet_tumours[slice_show], 'FDG', cmap = \"hot\")\n",
    "        subplot_([1,3,2], mask_np[slice_show], 'mask', cmap=\"hot\")\n",
    "        subplot_([1,3,3], tumours[slice_show]/tumours.max() + backgrounds[slice_show], 'tumour+background', cmap=\"hot\")\n",
    "        plt.show() \"\"\"\n",
    "    print(\"The resolution is: \", vol[\"res\"], \"mm. The dimensions are: \", pet_tumours.shape, \"voxels.\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sirf.Utilities import examples_data_path\n",
    "import sirf.STIR as pet\n",
    "pet.set_verbosity(0)\n",
    "pet.AcquisitionData.set_storage_scheme(\"memory\")\n",
    "pet.MessageRedirector(info=None, warn=None, errr=None)\n",
    "import numpy as np\n",
    "\n",
    "for path in [\"brainweb_3D/subject_04.FDG.MR.ground_truth.tumours.npy\", \"brainweb_3D/subject_04.Amyloid.MR.ground_truth.tumours.npy\"]:\n",
    "    gt_tumour_dict = np.load(path, allow_pickle=True).item()\n",
    "\n",
    "    mMR_template_sino = examples_data_path('PET') + \"/mMR/mMR_template_span11.hs\"\n",
    "    mMR_template = pet.AcquisitionData(mMR_template_sino)\n",
    "    image = mMR_template.create_uniform_image(1.0)\n",
    "    image_hr = image.zoom_image(\n",
    "                    zooms=(2.03125/1, 2.08626/1, 2.08626/1),\n",
    "                    offsets_in_mm=(0., 0., 0.),\n",
    "                    size=(180, 256, 256))\n",
    "    image_hr = image_hr.move_to_scanner_centre(mMR_template)\n",
    "    image_hr.fill(gt_tumour_dict[\"PET_tumours\"])\n",
    "    image_lr = image_hr.zoom_image(\n",
    "                zooms=(1/2, 1/2, 1/2),\n",
    "                offsets_in_mm=(0., 0., 0.),\n",
    "                size=(90,128,128))\n",
    "    image_lr = image_lr.move_to_scanner_centre(mMR_template)\n",
    "\n",
    "    print(\"Image original shape: \", image.shape, \". Voxel dimensions (mm): \", image.voxel_sizes())\n",
    "    print(\"Image high resolution shape: \", image_hr.shape, \". Voxel dimensions (mm): \", image_hr.voxel_sizes())\n",
    "    print(\"Image low resolution shape: \", image_lr.shape, \". Voxel dimensions (mm): \", image_lr.voxel_sizes())\n",
    "\n",
    "    # HIGH RESOLUTION\n",
    "    am = pet.AcquisitionModelUsingParallelproj()\n",
    "    # Set up sensitivity due to attenuation\n",
    "    uMap = image_hr.clone().fill(gt_tumour_dict[\"uMap\"])\n",
    "    asm_attn = pet.AcquisitionSensitivityModel(uMap, am)\n",
    "    asm_attn.set_up(mMR_template)\n",
    "    bin_eff_hr = mMR_template.get_uniform_copy(1.)\n",
    "    print('applying attenuation (please wait, may take a while)...')\n",
    "    asm_attn.unnormalise(bin_eff_hr)\n",
    "    asm = pet.AcquisitionSensitivityModel(bin_eff_hr)\n",
    "    am.set_acquisition_sensitivity(asm)\n",
    "    am.set_up(mMR_template, image_hr)\n",
    "    clean_measurements_hr = am.forward(image_hr)\n",
    "\n",
    "    \n",
    "    # HIGH RESOLUTION\n",
    "    am = pet.AcquisitionModelUsingParallelproj()\n",
    "    # Set up sensitivity due to attenuation\n",
    "    uMap = image_hr.clone().fill(gt_tumour_dict[\"uMap\"])\n",
    "    asm_attn = pet.AcquisitionSensitivityModel(uMap, am)\n",
    "    asm_attn.set_up(mMR_template)\n",
    "    bin_eff_hr = mMR_template.get_uniform_copy(1.)\n",
    "    print('applying attenuation (please wait, may take a while)...')\n",
    "    asm_attn.unnormalise(bin_eff_hr)\n",
    "    asm = pet.AcquisitionSensitivityModel(bin_eff_hr)\n",
    "    am.set_acquisition_sensitivity(asm)\n",
    "    am.set_up(mMR_template, image_lr)\n",
    "    clean_measurements_lr = am.forward(image_lr)\n",
    "\n",
    "    print(\"Rescaling for change image size with gt efficiencies: \", clean_measurements_lr.sum()/clean_measurements_hr.sum())\n",
    "\n",
    "    name = \"brainweb_3D/\"\n",
    "    if \"Amyloid\" in path:\n",
    "        name += \"Amyloid_\"\n",
    "    else:\n",
    "        name += \"FDG_\"\n",
    "\n",
    "    image_hr.write(name+\"PET_hr\")\n",
    "    clean_measurements_hr.write(name+\"clean_measurements_hr\")\n",
    "    bin_eff_hr.write(name+\"bin_eff_hr\")\n",
    "\n",
    "    image_lr.write(name+\"PET_lr\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import torch\n",
    "pet_old = torch.load(\"path_to/src/brainweb_2d/clean/clean_test.pt\")[\"reference\"].detach().cpu().numpy()\n",
    "image_fdg = pet.ImageData(\"path_to/src/sirf/brainweb_3D/FDG_PET_lr.hv\")\n",
    "image_amyl = pet.ImageData(\"path_to/src/sirf/brainweb_3D/Amyloid_PET_lr.hv\")\n",
    "fig, ax = plt.subplots(1,3,figsize=(15,5))\n",
    "slice_show = 30\n",
    "# array manipulation .swapaxes(0,1)[::-1,::-1]\n",
    "fig.colorbar(ax[0].imshow(pet_old[slice_show,0,...].swapaxes(0,1)[::-1,::-1]))\n",
    "ax[0].set_title(\"Original\")\n",
    "fig.colorbar(ax[1].imshow(image_fdg.as_array()[slice_show,...]))\n",
    "ax[1].set_title(\"FDG\")\n",
    "fig.colorbar(ax[2].imshow(image_amyl.as_array()[slice_show,...]))\n",
    "ax[2].set_title(\"Amyloid\")\n",
    "pet_slice = pet_old[slice_show,0,...]\n",
    "plt.show()\n",
    "print(np.sum(pet_slice.swapaxes(0,1)[::-1,::-1].swapaxes(0,1)[::-1,::-1]-pet_slice)) \n",
    "\n",
    "# Validate that the old processed data is similar orientation/space as new data\n",
    "import torch\n",
    "pet_old = torch.load(\"path_to/src/brainweb_2d/clean/clean_test.pt\")[\"reference\"].detach().cpu().numpy()\n",
    "image_fdg = pet.ImageData(\"path_to/src/sirf/brainweb_3D/FDG_PET_hr.hv\")\n",
    "image_amyl = pet.ImageData(\"path_to/src/sirf/brainweb_3D/Amyloid_PET_hr.hv\")\n",
    "fig, ax = plt.subplots(1,3,figsize=(15,5))\n",
    "# array manipulation .swapaxes(0,1)[::-1,::-1]\n",
    "fig.colorbar(ax[0].imshow(pet_old[slice_show,0,...].swapaxes(0,1)[::-1,::-1]))\n",
    "ax[0].set_title(\"Original\")\n",
    "fig.colorbar(ax[1].imshow(image_fdg.as_array()[2*slice_show,...]))\n",
    "ax[1].set_title(\"FDG\")\n",
    "fig.colorbar(ax[2].imshow(image_amyl.as_array()[2*slice_show,...]))\n",
    "ax[2].set_title(\"Amyloid\")\n",
    "pet_slice = pet_old[slice_show,0,...]\n",
    "plt.show()\n",
    "print(np.sum(pet_slice.swapaxes(0,1)[::-1,::-1].swapaxes(0,1)[::-1,::-1]-pet_slice))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for path in [\"brainweb_3D/subject_04.FDG.MR.ground_truth.tumours.npy\", \"brainweb_3D/subject_04.Amyloid.MR.ground_truth.tumours.npy\"]:\n",
    "    gt_tumour_dict = np.load(path, allow_pickle=True).item()\n",
    "    tumour_dict = gt_tumour_dict[\"tumours_dict\"]\n",
    "    name = \"brainweb_3D/\"\n",
    "    if \"Amyloid\" in path:\n",
    "        name += \"Amyloid_\"\n",
    "    else:\n",
    "        name += \"FDG_\"\n",
    "    for i in range(5):\n",
    "        for masks in [\"tumour_mask\", \"background_mask\"]:\n",
    "            mask = tumour_dict[masks+\"_\"+str(i)]\n",
    "            mMR_template_sino = examples_data_path('PET') + \"/mMR/mMR_template_span11.hs\"\n",
    "            mMR_template = pet.AcquisitionData(mMR_template_sino)\n",
    "            image = mMR_template.create_uniform_image(1.0)\n",
    "            image_hr = image.zoom_image(\n",
    "                            zooms=(2.03125/1, 2.08626/1, 2.08626/1),\n",
    "                            offsets_in_mm=(0., 0., 0.),\n",
    "                            size=(180, 256, 256))\n",
    "            image_hr = image_hr.move_to_scanner_centre(mMR_template)\n",
    "            image_hr.fill(mask)\n",
    "            image_lr = image_hr.zoom_image(\n",
    "                        zooms=(1/2, 1/2, 1/2),\n",
    "                        offsets_in_mm=(0., 0., 0.),\n",
    "                        size=(90,128,128))\n",
    "            image_lr = image_lr.move_to_scanner_centre(mMR_template)\n",
    "            image_lr.write(name+masks+\"_\"+str(i))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sirf.STIR as pet\n",
    "import torch\n",
    "tumours = []\n",
    "backgrounds = []\n",
    "for i in range(5):\n",
    "    tumour = pet.ImageData((\"brainweb_3D/Amyloid_tumour_mask_\"+str(i))+\".hv\").as_array()\n",
    "    background = pet.ImageData((\"brainweb_3D/Amyloid_background_mask_\"+str(i))+\".hv\").as_array()\n",
    "    tumours.append(torch.from_numpy(tumour))\n",
    "    backgrounds.append(torch.from_numpy(background))\n",
    "tumours = torch.stack(tumours)\n",
    "backgrounds = torch.stack(backgrounds)\n",
    "torch.save(tumours, \"brainweb_3D/tumours.pt\")\n",
    "torch.save(backgrounds, \"brainweb_3D/backgrounds.pt\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Uncomment to generate noisy data\n",
    "import sirf.STIR as pet\n",
    "pet.set_verbosity(0)\n",
    "from sirf.Utilities import examples_data_path\n",
    "pet.AcquisitionData.set_storage_scheme(\"memory\")\n",
    "pet.MessageRedirector(info=None, warn=None, errr=None)\n",
    "import os\n",
    "os.makedirs(\"brainweb_3D/noisy\", exist_ok=True)\n",
    "import numpy as np\n",
    "num_realisations = 5\n",
    "np.random.seed(42)\n",
    "for tracer in [\"Amyloid\", \"FDG\"]:\n",
    "    for num_counts in [4e7]:\n",
    "        for realisation in range(num_realisations):\n",
    "            name = \"brainweb_3D/noisy/\" + tracer + \"_noisy_measurements_\" + \"{:.0e}\".format(num_counts) + \"_\" + str(realisation)\n",
    "            clean_measurements = pet.AcquisitionData(\"brainweb_3D/\" + tracer + \"_clean_measurements_hr.hs\")\n",
    "            clean_measurements_arr = np.abs(clean_measurements.as_array())\n",
    "            print(\"Noise rescaling factor: \", num_counts/clean_measurements_arr.sum())\n",
    "            clean_measurements_arr = clean_measurements_arr/clean_measurements_arr.sum()*num_counts\n",
    "            noisy_measurements_arr = np.random.poisson(clean_measurements_arr).astype('float32')\n",
    "            noisy_measurements = clean_measurements.get_uniform_copy()\n",
    "            noisy_measurements.fill(noisy_measurements_arr)\n",
    "            noisy_measurements.write(name)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get 1 epoch OSEM images for noises and realisations\n",
    "import sirf.STIR as pet\n",
    "import numpy as np\n",
    "pet.set_verbosity(0)\n",
    "from sirf.Utilities import examples_data_path\n",
    "pet.AcquisitionData.set_storage_scheme(\"memory\")\n",
    "pet.MessageRedirector(info=None, warn=None, errr=None)\n",
    "num_realisations = 5\n",
    "num_subsets = 28\n",
    "num_epochs = 1\n",
    "for tracer in [\"Amyloid\", \"FDG\"]:\n",
    "    image = pet.ImageData(f\"brainweb_3D/{tracer}_PET_lr.hv\")\n",
    "    bin_eff = pet.AcquisitionData(f\"brainweb_3D/{tracer}_bin_eff_hr.hs\")\n",
    "    data_template = bin_eff.get_uniform_copy(1.)\n",
    "    image_template = image.get_uniform_copy(1.)\n",
    "    am = pet.AcquisitionModelUsingParallelproj()\n",
    "    asm = pet.AcquisitionSensitivityModel(bin_eff)\n",
    "    am.set_acquisition_sensitivity(asm)\n",
    "    am.set_up(data_template, image.get_uniform_copy(1))\n",
    "    sens_image = am.backward(data_template)\n",
    "    sens_image.write(f\"brainweb_3D/{tracer}_sensitivity_image.hv\")\n",
    "    sens_image_np = sens_image.as_array()\n",
    "    views = data_template.shape[2]\n",
    "    acquisition_models_sirf = []\n",
    "    for i in range(num_subsets):\n",
    "        subset_idxs = list(range(views))[i:][::num_subsets]\n",
    "        bin_eff_subset = bin_eff.get_subset(subset_idxs)\n",
    "        data_template_subset = bin_eff_subset.get_uniform_copy(1.)\n",
    "        # SET UP THE ACQUISITION MODEL\n",
    "        acquisition_model = pet.AcquisitionModelUsingParallelproj()\n",
    "        sensitivity_factors = pet.AcquisitionSensitivityModel(bin_eff_subset)\n",
    "        acquisition_model.set_acquisition_sensitivity(sensitivity_factors)\n",
    "        acquisition_model.set_up(data_template_subset, image)\n",
    "        acquisition_models_sirf.append(acquisition_model)\n",
    "        print(\"Acquisition model\", i+1, \"of\", num_subsets, \"set up\")\n",
    "    for num_counts in [4e7]:\n",
    "        for realisation in range(num_realisations):\n",
    "            name = \"brainweb_3D/noisy/\" + tracer + \"_noisy_measurements_\" + \"{:.0e}\".format(num_counts) + \"_\" + str(realisation)\n",
    "            noisy_measurements = pet.AcquisitionData(name + \".hs\")\n",
    "            objectives_sirf = []\n",
    "            for i in range(num_subsets):\n",
    "                subset_idxs = list(range(views))[i:][::num_subsets]\n",
    "                objective_sirf = pet.make_Poisson_loglikelihood(noisy_measurements.get_subset(subset_idxs), acq_model = acquisition_models_sirf[i])\n",
    "                objective_sirf.set_up(image)\n",
    "                objectives_sirf.append(objective_sirf)\n",
    "                print(\"Objective function\", i+1, \"of\", num_subsets, \"set up\")\n",
    "            x = image.get_uniform_copy(1)\n",
    "            views = noisy_measurements.shape[2]\n",
    "            for j in range(num_subsets):\n",
    "                bp_diff = objectives_sirf[j].get_gradient(x)\n",
    "                x = x + num_subsets * ((x+1e-6)/sens_image)*bp_diff\n",
    "                x_np = x.as_array()\n",
    "                x_np[sens_image_np==0] = 0\n",
    "                x_np[x_np<0] = 0\n",
    "                x.fill(x_np)\n",
    "            print(\"OSEM \", tracer, num_counts, realisation, \" done\")\n",
    "            x.write(\"brainweb_3D/noisy/\" + tracer + \"_osem_\" + \"{:.0e}\".format(num_counts) + \"_\" + str(realisation))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sirf.STIR as pet\n",
    "pet.set_verbosity(0)\n",
    "from sirf.Utilities import examples_data_path\n",
    "pet.AcquisitionData.set_storage_scheme(\"memory\")\n",
    "import numpy as np\n",
    "num_subsets = 28\n",
    "num_epochs = 10\n",
    "noise = \"4e+07\"\n",
    "tracer = \"Amyloid\"\n",
    "image = pet.ImageData(f\"brainweb_3D/{tracer}_PET_lr.hv\")\n",
    "measurements = pet.AcquisitionData(f\"brainweb_3D/noisy/{tracer}_noisy_measurements_{noise}_0.hs\")\n",
    "bin_eff = pet.AcquisitionData(f\"brainweb_3D/{tracer}_bin_eff_hr.hs\")\n",
    "\n",
    "am = pet.AcquisitionModelUsingParallelproj()\n",
    "asm = pet.AcquisitionSensitivityModel(bin_eff)\n",
    "am.set_acquisition_sensitivity(asm)\n",
    "am.set_up(measurements, image.get_uniform_copy(1))\n",
    "sens_image = am.backward(measurements.get_uniform_copy(1.))\n",
    "views = measurements.shape[2]\n",
    "noisy_measurements_subsets_sirf = []\n",
    "acquisition_models_sirf = []\n",
    "objectives_sirf = []\n",
    "for i in range(num_subsets):\n",
    "    subset_idxs = list(range(views))[i:][::num_subsets]\n",
    "    # Get subset of data\n",
    "    noisy_measurements_sirf = measurements.get_subset(subset_idxs)\n",
    "    bin_eff_subset = bin_eff.get_subset(subset_idxs)\n",
    "\n",
    "    # List noisy subsets\n",
    "    noisy_measurements_subsets_sirf.append(noisy_measurements_sirf)\n",
    "\n",
    "    # SET UP THE ACQUISITION MODEL\n",
    "    sensitivity_factors = pet.AcquisitionSensitivityModel(bin_eff_subset)\n",
    "    acquisition_model = pet.AcquisitionModelUsingParallelproj()\n",
    "    acquisition_model.set_acquisition_sensitivity(sensitivity_factors)\n",
    "    acquisition_model.set_up(noisy_measurements_subsets_sirf[i], image)\n",
    "    acquisition_models_sirf.append(acquisition_model)\n",
    "    objective_sirf = pet.make_Poisson_loglikelihood(noisy_measurements_subsets_sirf[i], acq_model = acquisition_models_sirf[i])\n",
    "    objective_sirf.set_up(image)\n",
    "    objectives_sirf.append(objective_sirf)\n",
    "    print(\"Acquisition model\", i+1, \"of\", num_subsets, \"set up\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x = image.get_uniform_copy(1)\n",
    "# INITIALISE THE RECONSTRUCTION OBJECT\n",
    "import matplotlib.pyplot as plt\n",
    "sens_image_np = sens_image.as_array()\n",
    "for i in range(10):\n",
    "    print(f\"Epoch {i+1}\")\n",
    "    for j in range(num_subsets):\n",
    "        bp_diff = objectives_sirf[j].get_gradient(x)\n",
    "        x = x + num_subsets * ((x+1e-6)/sens_image)*bp_diff\n",
    "        x_np = x.as_array()\n",
    "        x_np[sens_image_np==0] = 0\n",
    "        x_np[x_np<0] = 0\n",
    "        x.fill(x_np)\n",
    "    plt.imshow(x.as_array()[30])\n",
    "    plt.colorbar()\n",
    "    plt.show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.11"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
